package com.heiye.blog.ai.advisor.impl;

import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.google.common.collect.Lists;
import com.heiye.blog.ai.advisor.DashScopeAdvisor;
import com.heiye.blog.ai.domain.dos.ChatMessageDO;
import com.heiye.blog.ai.domain.mapper.ChatMessageMapper;
import com.heiye.blog.ai.model.dto.AIChatRequest;
import com.heiye.blog.ai.model.vo.AiChatReqVO;
import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Flux;

import java.util.Comparator;
import java.util.List;
import java.util.Objects;

/**
 * @author: heiye
 * @date: 2025/11/09 下午9:29
 * @version: v1.0.0
 * @description: 自定义百炼模型对话记忆 Advisor
 */
@Slf4j
public class CustomDashScopeChatMemoryAdvisor implements DashScopeAdvisor {

    private final ChatMessageMapper chatMessageMapper;
    private final AIChatRequest aiChatRequest;
    private final int limit;

    public CustomDashScopeChatMemoryAdvisor(ChatMessageMapper chatMessageMapper, AIChatRequest aiChatRequest, int limit) {
        this.chatMessageMapper = chatMessageMapper;
        this.aiChatRequest = aiChatRequest;
        this.limit = limit;
    }

    /**
     * 处理请求参数
     *
     * @param generationParam
     * @return
     */
    @Override
    public GenerationParam adviseRequest(GenerationParam generationParam) {
        log.info("## 自定义百炼模型聊天记忆 Advisor...");

        // 对话 UUID
        String chatUuid = aiChatRequest.getChatId();

        // 查询数据库拉取最新的聊天消息
        List<ChatMessageDO> messages = chatMessageMapper.selectList(Wrappers.<ChatMessageDO>lambdaQuery()
                // 查询指定对话 UUID 下的聊天记录
                .eq(ChatMessageDO::getChatUuid, chatUuid)
                // 查询最新的消息
                .orderByDesc(ChatMessageDO::getCreateTime)
                // 仅查询 LIMIT 条
                .last(String.format("LIMIT %d", limit)));

        // 按发布时间升序排列
        List<ChatMessageDO> sortedMessages = messages.stream()
                // 升序排列
                .sorted(Comparator.comparing(ChatMessageDO::getCreateTime))
                .toList();

        // 所有消息
        List<Message> messageList = Lists.newArrayList();

        // 将数据库记录转换为对应类型消息
        for (ChatMessageDO chatMessageDO : sortedMessages) {
            // 消息类型
            // 用户消息
            if (Objects.equals(chatMessageDO.getRole(), Role.USER.getValue())) {

                // 构建调用百炼模型的 message
                Message message = Message.builder()
                        .role(Role.USER.getValue())
                        .content(chatMessageDO.getContent())
                        .build();

                messageList.add(message);
            }
            // AI 助手消息
            else if (Objects.equals(chatMessageDO.getRole(), Role.ASSISTANT.getValue())) {

                Message message = Message.builder()
                        .role(Role.ASSISTANT.getValue())
                        .content(chatMessageDO.getContent())
                        .build();

                messageList.add(message);
            }
        }

        // 除了记忆消息，还需要添加当前用户消息
        messageList.addAll(generationParam.getMessages());

        // 设置消息
        generationParam.setMessages(messageList);

        // 返回新的GenerationParam
        return generationParam;
    }

    /**
     * 拦截处理
     *
     * @param fluxResult
     * @return
     */
    @Override
    public Flux<GenerationResult> adviseStream(Flux<GenerationResult> fluxResult) {
        // 对话记忆无需处理流数据
        return fluxResult;
    }

    /**
     * 执行顺序
     */
    @Override
    public int getOrder() {
        // order 值越小，越先执行
        return 2;
    }
}
